'''
Created on Oct 17, 2009

GenBank query code based on genbankdownload by Simon Greenhill

@author: ian
'''

import urllib, sys, os
from xml.etree.ElementTree import parse
import xml.etree.ElementTree

_email = 'imcnaney@gmail.com'
_tool = 'fetch-genbank'
_gb_url = 'http://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi?'
_db = 'nucleotide'
_data_path = 'c:/Users/ian/genomics/sequences/'

_scores = {'A': {'A':3, 'C':-1, 'T':-1, 'G': -1},
           'C': {'A':-1, 'C':3, 'T':-1, 'G': -1},
           'T': {'A':-1, 'C':-1, 'T':3, 'G': -1},
           'G': {'A':-1, 'C':-1, 'T':-1, 'G': 3}}

def cache_genbank(dom, filename):
    out = open(filename, 'w')
    out.write(xml.etree.ElementTree.tostring(dom.getroot()))
    out.flush()
    out.close()

def load_sequence(accession):
    """
    Load a nucleotide sequence into memory
    If it's on disk then load from there, else retrieve from genbank
    
    :param accession: the accession number from genbank
    
    :return: the ElementTree of the genbank xml
    """
    filename = _data_path + accession + '.xml'
    
    if(os.access(filename, os.R_OK)):
        #print 'Loading', accession, 'from disk.'
        dom = parse(filename)
    else:
        #print 'Unable to find', accession, 'on disk.  Fetching from GenBank.'
        dom = parse(query_genbank(accession))
        cache_genbank(dom, filename)
    
    return dom

def query_genbank(accession):
    """
    Fetch a sequence from GenBank by accession number

    :param accession: the accession number

    :return: the result of the genbank query
    :rtype: stream

    """

    params = {
              'db': 'nucleotide',
              'email': _email,
              'id': accession,
              'rettype': 'xml',
              'tool': _tool
              }

    fullurl = _gb_url + urllib.urlencode(params)
    return urllib.urlopen(fullurl)
    
def walk(sequence, scores, verbose):
    sum = 0
    for c in sequence:
        sum += scores[c]
        if(verbose):
            print c, scores[c], sum
        else:
            print sum

if __name__ == '__main__':
    from optparse import OptionParser
    
    parser = OptionParser(usage="usage: %prog [-n x] [-v] accession_number")
    parser.add_option("-n", "--nucleotide", dest="nucleotide",
                      help="which nucleotide is scored as a 3")
    parser.add_option("-v", "--verbose", action="store_true", dest="verbose",
                      default=False, help="produce verbose output")
    options, args = parser.parse_args()
    
    if(options.nucleotide not in _scores.iterkeys()):
        print "Error -", options.nucleotide, "is not a valid nucleotide"
        sys.exit(1)

    try:
        accession = args[0]
    except IndexError:
        parser.print_help()
        sys.exit(1)

    dom = load_sequence(accession)

    sequences = dom.getroot().findall('.//IUPACna')
    if(len(sequences) > 1):
        print "ALERT!  Multiple sequences found."
    for s in sequences:
        walk(s.text, _scores[options.nucleotide], options.verbose)

